# -*- coding: utf-8 -*-
"""
    Project name：code_potentialflow
# -------------------------------
    File name：check_SC_process.py
    Created on：2025/4/11 15:52
    Author：(input)
    Description:
"""
import numpy as np
from scipy.special import ellipk, ellipj
from scipy.optimize import root_scalar, minimize


class EllipticSCTransform:
    def __init__(self, width, height):
        """初始化 SC 变换，输入矩形宽度和高度。"""
        self.a = width / 2  # 矩形宽度的一半
        self.b = height  # 矩形高度
        self.k = None  # 模数 k
        self.k_prime = None  # 补模数 k'
        self.C = None  # SC 变换系数

    def solve_k_and_k_prime(self):
        """通过 b*K(k) - a*K(k') = 0 求解模数 k 和补模数 k'。"""

        def equation(k):
            k_prime = np.sqrt(1 - k ** 2)
            K_k = ellipk(k ** 2)
            K_k_prime = ellipk(k_prime ** 2)
            return self.b * K_k - self.a * K_k_prime

        sol = root_scalar(equation, bracket=[1e-6, 1 - 1e-6], method='bisect')
        if sol.converged:
            self.k = sol.root
            self.k_prime = np.sqrt(1 - self.k ** 2)
        else:
            raise RuntimeError("求解 k 和 k' 未收敛")

    def verify_k_and_k_prime(self):
        """验证 k^2 + k'^2 = 1。"""
        if self.k is None or self.k_prime is None:
            raise ValueError("请先求解 k 和 k'")
        sum_of_squares = self.k ** 2 + self.k_prime ** 2
        if not np.isclose(sum_of_squares, 1, rtol=1e-10):
            raise ValueError(f"验证失败：k^2 + k'^2 = {sum_of_squares}")
        print("验证通过：k^2 + k'^2 = 1")

    def compute_C(self):
        """通过 b = C*K(k') 计算 SC 变换系数 C。"""
        if self.k_prime is None:
            raise ValueError("请先求解 k 和 k'")
        K_k_prime = ellipk(self.k_prime ** 2)
        self.C = self.b / K_k_prime
        return self.C

    def verify_corner_mappings(self):
        """验证矩形四个角点的映射。"""
        if self.k is None or self.k_prime is None or self.C is None:
            raise ValueError("请先求解 k, k' 和 C")

        K_k = ellipk(self.k ** 2)
        K_k_prime = ellipk(self.k_prime ** 2)
        print(f"big K:{K_k}")
        print(f"big K':{K_k_prime}")

        point1 = -self.C * K_k
        point2 = self.C * K_k
        point3 = self.C * K_k + 1j * self.C * K_k_prime
        point4 = -self.C * K_k + 1j * self.C * K_k_prime

        expected_point1 = -self.a
        expected_point2 = self.a
        expected_point3 = self.a + 1j * self.b
        expected_point4 = -self.a + 1j * self.b

        print("角点映射验证：")
        print(f"点 1: -C*K(k) = {point1}, 预期: {expected_point1}")
        print(f"点 2: C*K(k) = {point2}, 预期: {expected_point2}")
        print(f"点 3: C*K(k) + iC*K(k') = {point3}, 预期: {expected_point3}")
        print(f"点 4: -C*K(k) + iC*K(k') = {point4}, 预期: {expected_point4}")

        if not all(np.isclose([point1, point2, point3, point4],
                              [expected_point1, expected_point2, expected_point3, expected_point4],
                              rtol=1e-10)):
            raise ValueError("角点映射验证失败")
        print("所有角点映射验证通过！")

    def compute_jacobi_sn(self, k, C, u):
        """
        计算 Jacobi 椭圆函数 sn(u, k) 并乘以系数 C，支持复数 u。
        参数：
            k: 模数 (0 <= k <= 1)
            C: 系数（标量）
            u: 复数坐标（u = X + iY）
        返回：
            z:  sn(u/c, k)，复数值
        """
        if not (0 <= k <= 1):
            raise ValueError(f"模数 k 必须在 [0, 1]，得到 {k}")
        if not isinstance(u, (int, float, complex)):
            raise ValueError(f"u 必须是标量或复数，得到 {type(u)}")
        if not isinstance(C, (int, float)):
            raise ValueError(f"C 必须是标量，得到 {type(C)}")

        try:
            # 分解 u = X + iY
            X = float(u.real)
            Y = float(u.imag)

            # 计算 k' 和 m = k^2, m' = k'^2
            k_prime = np.sqrt(1 - k ** 2)
            m = k ** 2
            m_prime = k_prime ** 2

            # 计算椭圆函数值
            sn_X_k, cn_X_k, dn_X_k, _ = ellipj(X / C, m)
            sn_Y_kp, cn_Y_kp, dn_Y_kp, _ = ellipj(Y / C, m_prime)


            # 按照公式计算 sn(X + iY, k)
            num = sn_X_k * dn_Y_kp + 1j * sn_Y_kp * cn_Y_kp * cn_X_k * dn_X_k
            den = cn_Y_kp ** 2 + k ** 2 * sn_X_k ** 2 * sn_Y_kp ** 2

            # 检查分母，避免除零
            if abs(den) < 1e-10:
                raise ValueError("分母接近零，可能导致数值不稳定")

            z = (num / den)
            return z
        except Exception as e:
            raise ValueError(f"计算 Jacobi 椭圆函数失败：{e}")

    def upper_half_to_rectangle(self, u):
        """将上半平面坐标 u 映射到矩形平面坐标 z。"""
        if self.C is None or self.k is None:
            raise ValueError("请先求解 C 和 k")

        sn, _, _, _ = ellipj(float(u) / float(self.C), float(self.k ** 2))
        z = self.C * sn
        return z

    def rectangle_to_upper_half(self, z):
        """将矩形平面坐标 z 映射回上半平面坐标 u。"""
        if self.C is None or self.k is None:
            raise ValueError("请先求解 C 和 k")

        # 检查 z 是否在矩形范围内
        if not (-self.a <= z.real <= self.a and 0 <= z.imag <= self.b):
            raise ValueError(f"坐标 {z} 超出矩形范围 [-{self.a}, {self.a}] x [0, {self.b}]")

        # 目标函数：最小化 |C * sn(u, k) - z|
        def objective(u):
            try:
                u_complex = complex(u[0], u[1])
                sn, _, _, _ = ellipj(u_complex / float(self.C), float(self.k ** 2))
                diff = self.C * sn - z
                return diff.real ** 2 + diff.imag ** 2
            except Exception as e:
                print(f"目标函数错误：u = {u}, 错误信息：{e}")
                return np.inf  # 返回无穷大以避免优化失败

        # 初始猜测：基于 z 的实部和虚部
        u_guess = np.array([z.real / self.C, z.imag / self.C])
        result = minimize(objective, u_guess, method='Powell')  # 使用 Powell 方法提高鲁棒性

        if result.success:
            u = complex(result.x[0], result.x[1])
            # 确保 u 在上半平面（虚部非负）
            if u.imag < -1e-10:  # 允许小误差
                u = u.conjugate()
            return u
        else:
            raise RuntimeError(f"求解上半平面坐标 u 未收敛：{result.message}")

    def verify_mapping(self, u):
        """验证双向映射一致性：u → z → u'。"""
        z = self.upper_half_to_rectangle(u)
        u_prime = self.rectangle_to_upper_half(z)
        print(f"原始 u: {u}, 映射到矩形 z: {z}, 逆映射回 u': {u_prime}")
        if not np.isclose(u.real, u_prime.real, rtol=1e-5) or not np.isclose(u.imag, u_prime.imag, rtol=1e-5):
            print("警告：映射验证未完全通过")
        else:
            print("映射验证通过！")
        return z, u_prime

    def solve(self):
        """执行完整求解流程。"""
        print("求解模数 k 和补模数 k'...")
        self.solve_k_and_k_prime()
        print(f"模数 k: {self.k}")
        print(f"补模数 k': {self.k_prime}")

        print("验证 k 和 k'...")
        self.verify_k_and_k_prime()

        print("计算 SC 变换系数 C...")
        self.C = self.compute_C()
        print(f"系数 C: {self.C}")

        print("验证矩形角点映射...")
        self.verify_corner_mappings()

        return self.C, self.k, self.k_prime


# 示例用法
if __name__ == "__main__":
    # 定义矩形尺寸
    width = 1.5
    height = 5.0
    transformer = EllipticSCTransform(width, height)

    # 求解参数
    C, k, k_prime = transformer.solve()

    # 测试 Jacobi 椭圆函数
    print("\n测试 Jacobi 椭圆函数计算：")
    # 原始点和复数坐标
    point_data = [
        {"label": "A", "u": -0.75},
        {"label": "B", "u": -0.75 + 5.0j},
        {"label": "C", "u": -0.25 + 5.0j},
        {"label": "D", "u": 0.75 + 5.0j},
        {"label": "E", "u": 0.75 + 0.5j},
        {"label": "F", "u": 0.75},
        {"label": "J", "u": 0.25 + 5.0j}
    ]

    # 添加放缩后的 u'，并计算变换后的 z
    for entry in point_data:
        u = entry["u"]
        up = u / C
        z = transformer.compute_jacobi_sn(k, C, u)
        entry["u_scaled"] = up
        entry["z_mapped"] = z

    # 打印结果
    print("\n角点、控制点的放缩、变换：")
    print(f"k = {k:.6f}, C = {C:.6f}\n")

    for entry in point_data:
        print(f"点 {entry['label']}：")
        print(f"  原坐标 u     = {entry['u']}")
        print(f"  放缩坐标 u'  = {entry['u_scaled']}")
        print(f"  映射后 z     = {entry['z_mapped']}")
        print("-" * 40)


    # 测试坐标映射
    #u_test = 1.0  # 上半平面测试点
    #print(f"\n测试坐标映射 u = {u_test}")
    #z, u_prime = transformer.verify_mapping(u_test)

    # 测试矩形内的点
    #z_test = 1.0 + 1.0j
    #print(f"\n测试逆映射 z = {z_test}")
    #u = transformer.rectangle_to_upper_half(z_test)
    #z_prime = transformer.upper_half_to_rectangle(u)
    #print(f"矩形坐标 z: {z_test}, 映射到上半平面 u: {u}, 逆映射回 z': {z_prime}")
    #if not np.isclose(z_test.real, z_prime.real, rtol=1e-5) or not np.isclose(z_test.imag, z_prime.imag, rtol=1e-5):
        #print("逆映射验证失败！")
    #else:
        #print("逆映射验证通过！")



